#pragma once
#include "copy.h"
#include <iostream>

/**
 * Wrapper around cublasSgemv().
 * Performs Z = alpha*A*X + beta*Y. 
 *
 * @param _transA whether or not to transpose matrix A. 0 is no, 1 is yes
 * @param alpha scaling factor for A*X
 * @param A a m by n matrix to be scaled and multiplied
 * @param X a length n vector to be multiplied
 * @param beta a scaling factor for vector Y
 * @param Y 
 * @param Z Preallocated buffer for result
 */
void sgemv
(
 	int _transA,
	float alpha,
	Matrix A,
	Vector X,
	float beta,
	Vector Y,
	Vector Z
)
{
	bool is_col_A;
	int m, n, lda, incx, incy, size_y;
	char transA;

	is_col_A = (A.stride == 1);

    bool opt = !is_col_A && _transA;
	m = opt ? A.sub.length : A.length;
	n = opt ? A.length : A.sub.length;

	lda = is_col_A ? A.sub.stride : A.stride;
	incx = X.stride; 
	incy = Y.stride;

    if(_transA && !opt){
        if(m != X.length){
            std::cerr << "sgemv: ERROR: mismatch in demensions of A and X (" << m
                << " and " << X.length << " should be equal)" << std::endl;
        }

        if(n != Y.length){
            std::cerr << "sgemv: ERROR: mismatch in demensions of A and Y (" << n 
                << " and " << Y.length << " should be equal)" << std::endl;
        }
    } else {
        if(n != X.length){
            std::cerr << "sgemv: ERROR: mismatch in demensions of A and X (" << n
                << " and " << X.length << " should be equal)" << std::endl;
        }

        if(m != Y.length){
            std::cerr << "sgemv: ERROR: mismatch in demensions of A and Y (" << m 
                << " and " << Y.length << " should be equal)" << std::endl;
        }
    }

	if(m != n && !is_col_A && !_transA) { //we need to copy A
		//warn user
		std::cerr << "sgemv: WARNING: Passing in a non-square, row major matrix witout transposing has incurred a copy operation" << std::endl;
		int size_A = n * m * sizeof(float);
		float *A_data;
		cudaMalloc(&A_data, size_A);
		SubMatrix A_sub(n, m, A_data, 0);
		Matrix A_copy(m, 1, A_sub);
		copy(A, A_copy);
		A = A_copy; //huzzah for args passed by copy
		lda = A.sub.stride;
		is_col_A = 1;
	}

	size_y = (1 + (Y.length * incy)) * sizeof(float);
	cudaMemcpy(Z.data, Y.data, size_y, cudaMemcpyDeviceToDevice);


	transA = (_transA ^ is_col_A) ? 'n' : 't';

     // std::cout << "Matrix extents: (" << A.length << ", " << A.sub.length << ")" << std::endl; 
     // std::cout << "Matrix strides: (" << A.stride << ", " << A.sub.stride << ")" << std::endl; 
    
     // std::cout << "transA: " << transA << " m: " << m << " n: " << n << std::endl; 
     // std::cout << "A: " << A.sub.data + A.sub.offset << " lda: " << lda << " X: " << X.data << " incx: " << incx << std::endl; 
    // std::cout << " beta: " << beta << " Y: " << Z.data << " incy: " << incy << std::endl; 
    
	cublasSgemv(transA, m, n, alpha, A.sub.data + A.sub.offset, lda, X.data, incx, beta, Z.data, incy);

}

template<typename Seq>
void printSequence(char* name, Seq x) {
  std::cout << name << ", Length: " << x.size() << ", Data: " << x.data << std::endl;
}


/**
 * Wrapper around cublasSgemv().
 * Performs Z = alpha*A*X + beta*Y. 
 *
 * @param _transA whether or not to transpoase matrix A. 0 is no, 1 is yes
 * @param alpha scaling factor for A*X
 * @param A a m by n matrix to be scaled and multiplied
 * @param X a length n vector to be multiplied
 * @param beta a scaling factor for vector Y
 * @param Y 
 * @param Z Preallocated buffer for result
 */
void sgemv
(
	int _transA,
	float alpha,
	Matrix A,
	stored_sequence<float> X,
	float beta,
	stored_sequence<float> Y,
	stored_sequence<float> Z
)
{
  /* printSequence("X", X); */
  /* printSequence("Y", Y); */
  /* printSequence("Z", Z); */
	Vector x(X.length, 1, X.data);
	Vector y(Y.length, 1, Y.data);
	Vector z(Z.length, 1, Z.data);
	sgemv(_transA, alpha, A, x, beta, y, z);
}


